import joblib
import numpy as np
import pandas as pd
import sys
from pathlib import Path
from featurework import prepare_data
import xgboost as xgb

# 设置项目目录结构
SCRIPT_DIR = Path(__file__).resolve().parent
PROJECT_ROOT = SCRIPT_DIR.parents[1]
MODEL_DIR = PROJECT_ROOT / "model"
OUTPUT_DIR = PROJECT_ROOT / "output"
DATA_DIR = PROJECT_ROOT / "data"


def calibrate_predictions(predictions, method='rank'):
    """校准预测结果以提高排名准确性"""
    if method == 'rank':
        # 排名校准：增强极端值的差异
        sorted_pred = np.sort(predictions)
        n = len(sorted_pred)
        rank_weights = np.linspace(0.8, 1.2, n)

        # 创建排名权重映射
        rank_map = {}
        for i, val in enumerate(sorted_pred):
            rank_map[val] = rank_weights[i]

        # 应用排名校准
        calibrated = np.array([rank_map[p] * p for p in predictions])
        return calibrated

    elif method == 'extreme':
        # 极端值增强：增强顶部和底部预测
        q10 = np.percentile(predictions, 10)
        q90 = np.percentile(predictions, 90)

        calibrated = np.zeros_like(predictions)
        for i, p in enumerate(predictions):
            if p > q90:
                calibrated[i] = p * 1.1
            elif p < q10:
                calibrated[i] = p * 0.9
            else:
                calibrated[i] = p
        return calibrated

    else:
        return predictions


def predict_top_stocks():
    """预测涨跌幅最大和最小的股票"""
    print("=" * 50)
    print("开始股票涨跌幅预测")
    print("=" * 50)

    # 准备预测数据
    print("准备预测数据...")
    try:
        _, _, X_test, stock_codes = prepare_data(predict=True)
        print(f"预测数据形状: {X_test.shape}")
        print(f"股票数量: {len(stock_codes)}")
    except Exception as e:
        print(f"数据准备失败: {str(e)}")
        import traceback
        traceback.print_exc()
        return

    # 检查数据
    if X_test.empty:
        print("错误: 测试数据为空")
        return

    # 加载模型
    print("加载预测模型...")

    # 加载LightGBM模型
    lgb_model_files = list(MODEL_DIR.glob("lgb_model_fold*.pkl"))
    lgb_models = []
    if lgb_model_files:
        for file in lgb_model_files:
            try:
                model = joblib.load(file)
                lgb_models.append(model)
                print(f"已加载LightGBM模型: {file.name}")
            except Exception as e:
                print(f"加载LightGBM模型失败: {file}, {str(e)}")
    else:
        print("警告: 未找到LightGBM模型文件")

    # 加载XGBoost模型
    xgb_model_files = list(MODEL_DIR.glob("xgb_model_fold*.pkl"))
    xgb_models = []
    if xgb_model_files:
        for file in xgb_model_files:
            try:
                model = joblib.load(file)
                xgb_models.append(model)
                print(f"已加载XGBoost模型: {file.name}")
            except Exception as e:
                print(f"加载XGBoost模型失败: {file}, {str(e)}")
    else:
        print("警告: 未找到XGBoost模型文件")

    if not lgb_models and not xgb_models:
        print("错误: 没有可用的模型进行预测")
        return

    # 集成预测
    print("进行集成预测...")

    # LightGBM预测
    lgb_predictions = np.zeros(len(X_test))
    for i, model in enumerate(lgb_models):
        fold_pred = model.predict(X_test)
        lgb_predictions += fold_pred
        print(f"LightGBM模型 {i + 1} 预测完成")

    if lgb_models:
        lgb_predictions /= len(lgb_models)

    # XGBoost预测
    xgb_predictions = np.zeros(len(X_test))
    if xgb_models:
        dtest = xgb.DMatrix(X_test)
        for i, model in enumerate(xgb_models):
            fold_pred = model.predict(dtest)
            xgb_predictions += fold_pred
            print(f"XGBoost模型 {i + 1} 预测完成")
        xgb_predictions /= len(xgb_models)

    # 模型融合 (50% LightGBM + 50% XGBoost)
    if lgb_models and xgb_models:
        predictions = 0.50 * lgb_predictions + 0.50 * xgb_predictions
    elif lgb_models:
        predictions = lgb_predictions
    else:
        predictions = xgb_predictions

    # 校准预测结果
    print("校准预测结果...")
    calibrated_pred = calibrate_predictions(predictions, method='rank')

    # 创建结果DataFrame
    result_df = pd.DataFrame({
        '股票代码': stock_codes.values,
        '预测涨跌幅': calibrated_pred
    })

    # 按预测涨跌幅排序
    result_df = result_df.sort_values('预测涨跌幅', ascending=False)

    # 获取涨幅最大和最小的股票
    top10 = result_df.head(10)
    bottom10 = result_df.tail(10).sort_values('预测涨跌幅', ascending=True)

    print("\n涨幅最大的10支股票:")
    print(top10[['股票代码', '预测涨跌幅']].reset_index(drop=True))

    print("\n涨幅最小的10支股票:")
    print(bottom10[['股票代码', '预测涨跌幅']].reset_index(drop=True))

    # 创建最终结果
    final_result = pd.DataFrame({
        '涨幅最大股票代码': top10['股票代码'].tolist(),
        '涨幅最小股票代码': bottom10['股票代码'].tolist()
    })

    # 保存结果
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    output_path = OUTPUT_DIR / "result.csv"
    final_result.to_csv(output_path, index=False, encoding='utf-8')
    print(f"\n预测结果已保存至: {output_path}")

    print("=" * 50)
    print("预测任务完成!")
    print("=" * 50)


def main():
    """主函数"""
    # 确保必要的目录存在
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    MODEL_DIR.mkdir(parents=True, exist_ok=True)
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

    # 执行预测
    predict_top_stocks()


if __name__ == "__main__":
    main()